// Copyright (c) 2022 by Duguang.IO Inc. All Rights Reserved.
// Author: Ethan Liu
// Date: 2022-06-05 09:51:05

package engine

import (
	"context"
	"errors"
	"fmt"
	"github.com/sirupsen/logrus"
	"io"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/remotecommand"
	"path"
	"sync"
	"time"

	logger "jianmu-worker-kube/logging"

	"jianmu-worker-kube/engine/launcher"
	"jianmu-worker-kube/engine/podwatcher"

	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/kubernetes/scheme"
)

var errPodStopped = errors.New("pod has been stopped")

// Kubernetes Kubernetes引擎
type Kubernetes struct {
	client    kubernetes.Interface
	Client    kubernetes.Interface
	config    *rest.Config
	Config    *rest.Config
	watchers  *sync.Map
	launchers *sync.Map

	containerStartTimeout time.Duration
}

// New 返回一个新的Kubernetes引擎
func New(client kubernetes.Interface, config *rest.Config, containerStartTimeout time.Duration) *Kubernetes {
	if containerStartTimeout < time.Second {
		containerStartTimeout = time.Second
	}

	k := &Kubernetes{
		client:    client,
		config:    config,
		watchers:  &sync.Map{},
		launchers: &sync.Map{},

		containerStartTimeout: containerStartTimeout,
	}

	k.Client = client
	k.Config = config
	return k
}

// CreatePod 创建Pod
func (k *Kubernetes) CreatePod(ctx context.Context, unit *Unit) error {
	log := logger.FromContext(ctx).
		WithField("pod", unit.PodSpec.Name).
		WithField("namespace", unit.PodSpec.Namespace)

	// 创建cm
	// resultFile
	if unit.ResultFile != "" && path.Base(unit.ResultFile) != "" {
		unit.Runners[0].Envs["JM_RESULT_FILE"] = unit.ResultFile
	}

	labels := map[string]string{
		"executor":  "jianmu-k8s-worker",
		"creator":   "jianmu",
		"task-id":   unit.TaskInstanceId,
		"worker-id": unit.WorkerId,
	}

	var cm = &v1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:   unit.PodSpec.Name,
			Labels: labels,
		},
		Data: unit.Runners[0].Envs,
	}

	log.WithFields(logrus.Fields{
		"namespace": unit.PodSpec.Namespace,
	}).Debug("creating config map")

	_, err := k.client.CoreV1().ConfigMaps(unit.PodSpec.Namespace).
		Create(ctx, cm, metav1.CreateOptions{})

	if err != nil {
		log.WithError(err).Error("failed to create config map")
		unit.Log.Write([]byte(fmt.Sprintf("## 创建 configmap error: %s/%s %v\n", unit.PodSpec.Namespace, unit.PodSpec.Name, err)))
		return err
	} else {
		unit.Log.Write([]byte(fmt.Sprintf("## 创建 configmap succee: %s/%s\n", unit.PodSpec.Namespace, unit.PodSpec.Name)))
	}

	// 创建secret
	if unit.PullSecret != nil {
		pullSecret := toDockerConfigSecret(unit)
		_, err := k.client.CoreV1().Secrets(unit.PodSpec.Namespace).Create(ctx, pullSecret, metav1.CreateOptions{})
		if err != nil {
			log.WithError(err).Error("failed to create pull secret")
			return err
		}
		log.Trace("created pull secret")
	}

	// 创建pod
	pod := toPod(unit)
	pod.ObjectMeta.Labels = labels
	_, err = k.client.CoreV1().Pods(unit.PodSpec.Namespace).Create(ctx, pod, metav1.CreateOptions{})
	if err != nil {
		log.WithError(err).Error("failed to create pod")
		unit.Log.Write([]byte(fmt.Sprintf("##  创建pod失败: %s/%s %v\n", unit.PodSpec.Namespace, unit.PodSpec.Name, err)))
		return err
	} else {
		unit.Log.Write([]byte(fmt.Sprintf("## 创建pod成功: %s/%s image:%s \n", unit.PodSpec.Namespace, unit.PodSpec.Name, unit.Runners[0].Image)))
	}
	log.Trace("created pod")

	unit.Stop = make(chan struct{})
	return nil
}

func (k *Kubernetes) GetRunningPods(workerId, nameSpace string) ([]v1.Pod, error) {
	labelSelector := metav1.LabelSelector{
		MatchLabels: map[string]string{
			"executor":  "jianmu-k8s-worker",
			"creator":   "jianmu",
			"worker-id": workerId,
		},
	}

	listOptions := metav1.ListOptions{
		LabelSelector: labels.Set(labelSelector.MatchLabels).String(),
	}

	pods, err := k.client.CoreV1().Pods(nameSpace).List(context.TODO(), listOptions)
	if err != nil {
		logrus.WithError(err).
			WithField("labelSelector", labelSelector).
			WithField("nameSpace", nameSpace).
			Errorln("failed to list pods")
		return nil, err
	}
	return pods.Items, nil
}

func (k *Kubernetes) GetPodNodeName(podName, nameSpace string) (string, error) {
	pod, err := k.client.CoreV1().Pods(nameSpace).Get(context.TODO(), podName, metav1.GetOptions{})
	if err != nil {
		return "", err
	}
	return pod.Spec.NodeName, nil
}

// https://www.yfdou.com/archives/kuberneteszhi-kubectlexeczhi-ling-gong-zuo-yuan-li-shi-xian-copyhe-webshellyi-ji-filebrowser.html#Kubectl-exec-%E7%9A%84%E5%B7%A5%E4%BD%9C%E5%8E%9F%E7%90%86
func (k *Kubernetes) WaitPodRunning(ctx context.Context, podName, nameSpace string, out io.Writer) (running bool, err error, podSuccessed bool) {
	after := time.After(time.Second * 60 * 99)
	getNodeName := true
	events := make(map[string]*v1.Event)
	out.Write([]byte(fmt.Sprintf("## 等待 pod %s/%s 运行\n", nameSpace, podName)))

	for {
		select {
		case <-after:
			return false, errors.New("pod is not running"), false
		default:
			if getNodeName {
				if nodeName, err := k.GetPodNodeName(podName, nameSpace); err == nil {
					getNodeName = !getNodeName
					out.Write([]byte(fmt.Sprintf("## pod %s/%s 运行在[%s] node\n", nameSpace, podName, nodeName)))
				}
			}
			running, err, phase := GetPodStatus(ctx, k.client, podName, nameSpace, out, events)
			if phase == v1.PodSucceeded {
				return false, err, true
			}
			if err != nil {
				logger.FromContext(ctx).
					WithError(err).
					WithField("pod", podName).
					WithField("namespace", nameSpace).
					Error("failed to get pod status")
				return false, err, false
			} else if running {
				if getNodeName {
					if nodeName, err := k.GetPodNodeName(podName, nameSpace); err == nil {
						getNodeName = !getNodeName
						out.Write([]byte(fmt.Sprintf("## pod %s/%s 运行在[%s] node\n", nameSpace, podName, nodeName)))
					}
				}
				return true, nil, false
			}
			<-time.After(time.Second * 1)
		}
	}

}

func (k *Kubernetes) WaitContainerNameTermination(ctx context.Context, containerName, podName, nameSpace string) (exitCode int32, err error) {
	after := time.After(time.Second * 60 * 30)
	for {
		select {
		case <-after:
			return -1, errors.New("container run time out")
		default:
			running, err, exitCode := GetContainerStatus(ctx, k.client, containerName, podName, nameSpace)
			if running {
				<-time.After(time.Second * 5)
			} else {
				return exitCode, err
			}
		}
	}

}

// CopyFromContainer 从pod中读取文件
func (k *Kubernetes) CopyFromContainer(ctx context.Context, fileName, containerName, podName, nameSpace string, out io.Writer) error {
	return k.ExecCmd(ctx, []string{"/bin/sh", "-c", fmt.Sprintf("[[ -f %s ]] && cat %s || true", fileName, fileName)},
		containerName, podName, nameSpace, out)
}

func (k *Kubernetes) ExecCmd(ctx context.Context, cmd []string, containerName, podName, nameSpace string, out io.Writer) error {
	req := k.client.CoreV1().RESTClient().Post().
		Resource("pods").
		Name(podName).
		Namespace(nameSpace).
		SubResource("exec")

	req.VersionedParams(&v1.PodExecOptions{
		Container: containerName,
		Command:   cmd,
		Stdin:     false,
		Stdout:    true,
		Stderr:    false,
		TTY:       false,
	}, scheme.ParameterCodec)

	executor, err := remotecommand.NewSPDYExecutor(k.config, "POST", req.URL())
	if err != nil {
		logrus.WithError(err).Println("NewSPDYExecutor err")
		return err
	}

	// Stream
	err = executor.Stream(remotecommand.StreamOptions{
		Stdin:             nil,
		Stdout:            out,
		Stderr:            nil,
		TerminalSizeQueue: nil,
		Tty:               false,
	})

	return err
}

func (k *Kubernetes) DeleteCm(ctx context.Context, unit *Unit) error {
	log := logger.FromContext(ctx).
		WithField("cm", unit.PodSpec.Name).
		WithField("namespace", unit.PodSpec.Namespace)

	err := k.client.CoreV1().ConfigMaps(unit.PodSpec.Namespace).Delete(ctx, unit.PodSpec.Name, metav1.DeleteOptions{})
	if err != nil {
		log.Error("failed to delete config map")
	} else {
		log.Trace("deleted config map")
	}
	return err
}

// DeletePod 删除Pod
func (k *Kubernetes) DeletePod(ctx context.Context, unit *Unit) error {
	log := logger.FromContext(ctx).
		WithField("pod", unit.PodSpec.Name).
		WithField("namespace", unit.PodSpec.Namespace)
	if unit.PullSecret != nil {
		if err := k.client.CoreV1().Secrets(unit.PodSpec.Namespace).Delete(context.Background(), unit.PullSecret.Name, metav1.DeleteOptions{}); err != nil {
			log.WithError(err).Error("failed to delete pull secret")
		} else {
			log.Trace("deleted pull secret")
		}
	}

	close(unit.Stop)

	var isPodDeleted bool

	if err := k.client.CoreV1().Pods(unit.PodSpec.Namespace).Delete(context.Background(), unit.PodSpec.Name, metav1.DeleteOptions{}); err != nil {
		log.WithError(err).Error("failed to delete pod")
	} else {
		log.Trace("deleted pod")
		isPodDeleted = true
	}

	if _l, loaded := k.launchers.LoadAndDelete(unit.PodSpec.Name); loaded {
		l := _l.(*launcher.Launcher)
		l.Stop()
	}

	if w, loaded := k.watchers.LoadAndDelete(unit.PodSpec.Name); loaded {
		if isPodDeleted {
			watcher := w.(*podwatcher.PodWatcher)
			if err := watcher.WaitPodDeleted(unit.PodSpec.Name); err != nil && err != context.Canceled {
				log.WithError(err).Error("PodWatcher terminated with error")
			} else {
				log.Trace("PodWatcher terminated")
			}
		}
	}

	return nil
}

// Run 启动容器运行任务
func (k *Kubernetes) Run(ctx context.Context, unit *Unit, runner *Runner, output io.Writer) (state *State, err error) {

	podId := unit.PodSpec.Name
	podNamespace := unit.PodSpec.Namespace
	runnerName := runner.Name
	containerId := "build" //runner.ID
	containerImage := runner.Image
	containerPlaceholder := runner.Placeholder

	log := logger.FromContext(ctx).
		WithField("pod", podId).
		WithField("namespace", podNamespace).
		WithField("image", containerImage).
		WithField("placeholder", containerPlaceholder).
		WithField("container", containerId).
		WithField("step", runnerName)

	w, loaded := k.watchers.LoadOrStore(podId, &podwatcher.PodWatcher{})
	watcher := w.(*podwatcher.PodWatcher)
	if !loaded {
		watcher.Start(context.Background(), &podwatcher.KubernetesWatcher{
			PodNamespace: podNamespace,
			PodName:      podId,
			KubeClient:   k.client,
			Period:       20 * time.Second,
		})

		log.Trace("PodWatcher started")
	}

	err = watcher.AddContainer(runner.ID, runner.Placeholder, runner.Image)
	if err != nil {
		return
	}

	log.Debug("Engine: Starting unit")

	err = <-k.startContainer(ctx, unit, runner)
	if err != nil {
		return
	}

	chErrStart := make(chan error)
	go func() {
		chErrStart <- watcher.WaitContainerStart(containerId)
	}()

	select {
	case err = <-chErrStart:
	case <-time.After(k.containerStartTimeout):
		err = podwatcher.StartTimeoutContainerError{Container: containerId, Image: containerImage}
		log.WithError(err).Error("Engine: Container start timeout")
	case <-unit.Stop:
		return nil, errPodStopped
	}
	if err != nil {
		return
	}

	err = k.FetchLogs(ctx, unit, runner, output)
	if err != nil {
		return
	}

	type containerResult struct {
		code int
		err  error
	}

	chErrStop := make(chan containerResult)
	go func() {
		code, err := watcher.WaitContainerTerminated(containerId)
		chErrStop <- containerResult{code: code, err: err}
	}()

	select {
	case result := <-chErrStop:
		err = result.err
		if err != nil {
			return
		}

		state = &State{
			ExitCode:  result.code,
			Exited:    true,
			OOMKilled: false,
		}
	case <-unit.Stop:
		return nil, errPodStopped
	}

	return
}

func (k *Kubernetes) FetchLogs(ctx context.Context, unit *Unit, runner *Runner, output io.Writer) error {
	opts := &v1.PodLogOptions{
		Follow:    true,
		Container: runner.Name,
	}

	req := k.client.CoreV1().RESTClient().Get().
		Namespace(unit.PodSpec.Namespace).
		Name(unit.PodSpec.Name).
		Resource("pods").
		SubResource("log").
		VersionedParams(opts, scheme.ParameterCodec)

	readCloser, err := req.Stream(ctx)
	if err != nil {
		logger.FromContext(ctx).
			WithError(err).
			WithField("pod", unit.PodSpec.Name).
			WithField("namespace", unit.PodSpec.Namespace).
			WithField("container", runner.ID).
			WithField("runner", runner.Name).
			Error("failed to stream logs")
		return err
	}
	defer readCloser.Close()

	return cancellableCopy(ctx, output, readCloser)
}

func (k *Kubernetes) startContainer(ctx context.Context, unit *Unit, runner *Runner) <-chan error {
	podName := unit.PodSpec.Name
	podNamespace := unit.PodSpec.Namespace
	containerName := runner.Name
	containerImage := runner.Image

	_l, loaded := k.launchers.LoadOrStore(podName, launcher.New(podName, podNamespace, k.client, &unit.podUpdateMutex))
	l := _l.(*launcher.Launcher)
	if !loaded {
		l.Start(ctx)
	}

	return l.Launch(containerName, containerImage, runner.Envs)
}
